tests/mnle_test.py::test_mnle_api[none-mdn-vi]

Stack Traces | 0.048s run time

flow_model = 'mdn', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')
z_score_theta = 'none'

    @pytest.mark.parametrize(
        "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
    )
    @pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
    @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
    def test_mnle_api(
        flow_model: str,
        sampler,
        mcmc_params_fast: MCMCPosteriorParameters,
        z_score_theta: str,
    ):
        """Test MNLE API."""
        # Generate mixed data.
        num_simulations = 10
        theta = torch.rand(num_simulations, 2)
        x = torch.cat(
            (
                torch.rand(num_simulations, 1),
                torch.randint(0, 2, (num_simulations, 1)),
            ),
            dim=1,
        )
    
        # Train and infer.
        prior = BoxUniform(torch.zeros(2), torch.ones(2))
        x_o = x[0]
        # Build estimator manually.
        theta_embedding = FCEmbedding(2, 2)  # simple embedding net
        density_estimator = likelihood_nn(
            model="mnle",
            flow_model=flow_model,
            z_score_theta=z_score_theta,
            embedding_net=theta_embedding,
        )
        trainer = MNLE(density_estimator=density_estimator)
        trainer.append_simulations(theta, x).train(max_num_epochs=1)
    
        # Test different samplers.
>       posterior = trainer.build_posterior(prior=prior, sample_with=sampler)

tests/mnle_test.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nle/mnle.py:176: in build_posterior
    return super().build_posterior(
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.mnle.MNLE object at 0x7f8ccfd896f0>
estimator = MixedDensityEstimator(
  (net): NFlowsFlow(
    (net): Flow(
      (_transform): CompositeTransform(
        (_transfo...bias=True)
      (1): ReLU()
      (2): Linear(in_features=50, out_features=2, bias=True)
      (3): ReLU()
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'vi', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/mnle_test.py::test_mnle_api[independent-mdn-vi]

Stack Traces | 0.05s run time

flow_model = 'mdn', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')
z_score_theta = 'independent'

    @pytest.mark.parametrize(
        "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
    )
    @pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
    @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
    def test_mnle_api(
        flow_model: str,
        sampler,
        mcmc_params_fast: MCMCPosteriorParameters,
        z_score_theta: str,
    ):
        """Test MNLE API."""
        # Generate mixed data.
        num_simulations = 10
        theta = torch.rand(num_simulations, 2)
        x = torch.cat(
            (
                torch.rand(num_simulations, 1),
                torch.randint(0, 2, (num_simulations, 1)),
            ),
            dim=1,
        )
    
        # Train and infer.
        prior = BoxUniform(torch.zeros(2), torch.ones(2))
        x_o = x[0]
        # Build estimator manually.
        theta_embedding = FCEmbedding(2, 2)  # simple embedding net
        density_estimator = likelihood_nn(
            model="mnle",
            flow_model=flow_model,
            z_score_theta=z_score_theta,
            embedding_net=theta_embedding,
        )
        trainer = MNLE(density_estimator=density_estimator)
        trainer.append_simulations(theta, x).train(max_num_epochs=1)
    
        # Test different samplers.
>       posterior = trainer.build_posterior(prior=prior, sample_with=sampler)

tests/mnle_test.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nle/mnle.py:176: in build_posterior
    return super().build_posterior(
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.mnle.MNLE object at 0x7f8cd411d5a0>
estimator = MixedDensityEstimator(
  (net): NFlowsFlow(
    (net): Flow(
      (_transform): CompositeTransform(
        (_transfo...     (1): ReLU()
        (2): Linear(in_features=50, out_features=2, bias=True)
        (3): ReLU()
      )
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'vi', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/mnle_test.py::test_mnle_api[none-nsf-vi]

Stack Traces | 0.096s run time

flow_model = 'nsf', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')
z_score_theta = 'none'

    @pytest.mark.parametrize(
        "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
    )
    @pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
    @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
    def test_mnle_api(
        flow_model: str,
        sampler,
        mcmc_params_fast: MCMCPosteriorParameters,
        z_score_theta: str,
    ):
        """Test MNLE API."""
        # Generate mixed data.
        num_simulations = 10
        theta = torch.rand(num_simulations, 2)
        x = torch.cat(
            (
                torch.rand(num_simulations, 1),
                torch.randint(0, 2, (num_simulations, 1)),
            ),
            dim=1,
        )
    
        # Train and infer.
        prior = BoxUniform(torch.zeros(2), torch.ones(2))
        x_o = x[0]
        # Build estimator manually.
        theta_embedding = FCEmbedding(2, 2)  # simple embedding net
        density_estimator = likelihood_nn(
            model="mnle",
            flow_model=flow_model,
            z_score_theta=z_score_theta,
            embedding_net=theta_embedding,
        )
        trainer = MNLE(density_estimator=density_estimator)
        trainer.append_simulations(theta, x).train(max_num_epochs=1)
    
        # Test different samplers.
>       posterior = trainer.build_posterior(prior=prior, sample_with=sampler)

tests/mnle_test.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nle/mnle.py:176: in build_posterior
    return super().build_posterior(
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.mnle.MNLE object at 0x7f2e041f0b80>
estimator = MixedDensityEstimator(
  (net): NFlowsFlow(
    (net): Flow(
      (_transform): CompositeTransform(
        (_transfo...bias=True)
      (1): ReLU()
      (2): Linear(in_features=50, out_features=2, bias=True)
      (3): ReLU()
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'vi', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/mnle_test.py::test_mnle_api[independent-nsf-vi]

Stack Traces | 0.097s run time

flow_model = 'nsf', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')
z_score_theta = 'independent'

    @pytest.mark.parametrize(
        "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
    )
    @pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
    @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
    def test_mnle_api(
        flow_model: str,
        sampler,
        mcmc_params_fast: MCMCPosteriorParameters,
        z_score_theta: str,
    ):
        """Test MNLE API."""
        # Generate mixed data.
        num_simulations = 10
        theta = torch.rand(num_simulations, 2)
        x = torch.cat(
            (
                torch.rand(num_simulations, 1),
                torch.randint(0, 2, (num_simulations, 1)),
            ),
            dim=1,
        )
    
        # Train and infer.
        prior = BoxUniform(torch.zeros(2), torch.ones(2))
        x_o = x[0]
        # Build estimator manually.
        theta_embedding = FCEmbedding(2, 2)  # simple embedding net
        density_estimator = likelihood_nn(
            model="mnle",
            flow_model=flow_model,
            z_score_theta=z_score_theta,
            embedding_net=theta_embedding,
        )
        trainer = MNLE(density_estimator=density_estimator)
        trainer.append_simulations(theta, x).train(max_num_epochs=1)
    
        # Test different samplers.
>       posterior = trainer.build_posterior(prior=prior, sample_with=sampler)

tests/mnle_test.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nle/mnle.py:176: in build_posterior
    return super().build_posterior(
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.mnle.MNLE object at 0x7f8ccfdbba00>
estimator = MixedDensityEstimator(
  (net): NFlowsFlow(
    (net): Flow(
      (_transform): CompositeTransform(
        (_transfo...     (1): ReLU()
        (2): Linear(in_features=50, out_features=2, bias=True)
        (3): ReLU()
      )
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'vi', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/mnle_test.py::test_mnle_api[none-zuko_nsf-vi]

Stack Traces | 0.111s run time

flow_model = 'zuko_nsf', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')
z_score_theta = 'none'

    @pytest.mark.parametrize(
        "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
    )
    @pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
    @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
    def test_mnle_api(
        flow_model: str,
        sampler,
        mcmc_params_fast: MCMCPosteriorParameters,
        z_score_theta: str,
    ):
        """Test MNLE API."""
        # Generate mixed data.
        num_simulations = 10
        theta = torch.rand(num_simulations, 2)
        x = torch.cat(
            (
                torch.rand(num_simulations, 1),
                torch.randint(0, 2, (num_simulations, 1)),
            ),
            dim=1,
        )
    
        # Train and infer.
        prior = BoxUniform(torch.zeros(2), torch.ones(2))
        x_o = x[0]
        # Build estimator manually.
        theta_embedding = FCEmbedding(2, 2)  # simple embedding net
        density_estimator = likelihood_nn(
            model="mnle",
            flow_model=flow_model,
            z_score_theta=z_score_theta,
            embedding_net=theta_embedding,
        )
        trainer = MNLE(density_estimator=density_estimator)
        trainer.append_simulations(theta, x).train(max_num_epochs=1)
    
        # Test different samplers.
>       posterior = trainer.build_posterior(prior=prior, sample_with=sampler)

tests/mnle_test.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nle/mnle.py:176: in build_posterior
    return super().build_posterior(
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.mnle.MNLE object at 0x7f215ab33f40>
estimator = MixedDensityEstimator(
  (net): ZukoFlow(
    (net): Flow(
      (transform): LazyComposedTransform(
        (0): Unco...bias=True)
      (1): ReLU()
      (2): Linear(in_features=50, out_features=2, bias=True)
      (3): ReLU()
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'vi', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/mnle_test.py::test_mnle_api[independent-zuko_nsf-vi]

Stack Traces | 0.117s run time

flow_model = 'zuko_nsf', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')
z_score_theta = 'independent'

    @pytest.mark.parametrize(
        "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
    )
    @pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
    @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
    def test_mnle_api(
        flow_model: str,
        sampler,
        mcmc_params_fast: MCMCPosteriorParameters,
        z_score_theta: str,
    ):
        """Test MNLE API."""
        # Generate mixed data.
        num_simulations = 10
        theta = torch.rand(num_simulations, 2)
        x = torch.cat(
            (
                torch.rand(num_simulations, 1),
                torch.randint(0, 2, (num_simulations, 1)),
            ),
            dim=1,
        )
    
        # Train and infer.
        prior = BoxUniform(torch.zeros(2), torch.ones(2))
        x_o = x[0]
        # Build estimator manually.
        theta_embedding = FCEmbedding(2, 2)  # simple embedding net
        density_estimator = likelihood_nn(
            model="mnle",
            flow_model=flow_model,
            z_score_theta=z_score_theta,
            embedding_net=theta_embedding,
        )
        trainer = MNLE(density_estimator=density_estimator)
        trainer.append_simulations(theta, x).train(max_num_epochs=1)
    
        # Test different samplers.
>       posterior = trainer.build_posterior(prior=prior, sample_with=sampler)

tests/mnle_test.py:132: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nle/mnle.py:176: in build_posterior
    return super().build_posterior(
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.mnle.MNLE object at 0x7f2e0420cca0>
estimator = MixedDensityEstimator(
  (net): ZukoFlow(
    (net): Flow(
      (transform): LazyComposedTransform(
        (0): Unco...     (1): ReLU()
        (2): Linear(in_features=50, out_features=2, bias=True)
        (3): ReLU()
      )
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'vi', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/posterior_parameters_test.py::test_build_posterior_warns_on_conflicting_args[build_posterior_arguments1]

Flake rate in main: 33.33% (Passed 132 times, Failed 66 times)

Stack Traces | 0.041s run time

build_posterior_arguments = {'posterior_parameters': VIPosteriorParameters(q='maf', vi_method='fKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent'), 'vi_method': 'IW'}
get_inference = <sbi.inference.trainers.nre.nre_b.NRE_B object at 0x7eff5968e200>

    @pytest.mark.parametrize(
        "build_posterior_arguments",
        [
            dict(
                mcmc_method="slice_pymc",
                posterior_parameters=MCMCPosteriorParameters(method="hmc_pyro"),
            ),
            dict(
                vi_method="IW",
                posterior_parameters=VIPosteriorParameters(vi_method="fKL"),
            ),
        ],
    )
    def test_build_posterior_warns_on_conflicting_args(
        build_posterior_arguments, get_inference
    ):
        """
        Test that build_posterior raises a UserWarning on conflicting parameter
        combinations.
        """
        inference = get_inference
    
        with pytest.warns(UserWarning, match="ignored in favor of"):
>           inference.build_posterior(**build_posterior_arguments)

tests/posterior_parameters_test.py:198: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nre/nre_base.py:360: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nre.nre_b.NRE_B object at 0x7eff5968e200>
estimator = RatioEstimator(
  (net): ResidualNet(
    (initial_layer): Linear(in_features=6, out_features=50, bias=True)
    (bloc...Standardize()
    (1): Identity()
  )
  (embedding_net_x): Sequential(
    (0): Standardize()
    (1): Identity()
  )
)
prior = BoxUniform(Uniform(low: torch.Size([3]), high: torch.Size([3])), 1)
sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='fKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/posterior_parameters_test.py::test_build_posterior_works_on_default_args[build_posterior_arguments1]

Flake rate in main: 33.33% (Passed 132 times, Failed 66 times)

Stack Traces | 0.054s run time

build_posterior_arguments = {'posterior_parameters': VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')}
get_inference = <sbi.inference.trainers.nre.nre_b.NRE_B object at 0x7f0fd8b693f0>

    @pytest.mark.parametrize(
        "build_posterior_arguments",
        [
            pytest.param(
                dict(
                    posterior_parameters=MCMCPosteriorParameters(
                        method="slice_np_vectorized"
                    ),
                ),
            ),
            pytest.param(
                dict(
                    posterior_parameters=VIPosteriorParameters(vi_method="rKL"),
                ),
            ),
        ],
    )
    def test_build_posterior_works_on_default_args(
        build_posterior_arguments, get_inference
    ):
        """
        Test that build_posterior doesn't raise on default parameters.
        """
    
        inference = get_inference
>       inference.build_posterior(**build_posterior_arguments)

tests/posterior_parameters_test.py:226: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nre/nre_base.py:360: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nre.nre_b.NRE_B object at 0x7f0fd8b693f0>
estimator = RatioEstimator(
  (net): ResidualNet(
    (initial_layer): Linear(in_features=6, out_features=50, bias=True)
    (bloc...Standardize()
    (1): Identity()
  )
  (embedding_net_x): Sequential(
    (0): Standardize()
    (1): Identity()
  )
)
prior = BoxUniform(Uniform(low: torch.Size([3]), high: torch.Size([3])), 1)
sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/posterior_parameters_test.py::test_signature_consistency[VIPosteriorParameters-VIPosterior-skipped_fields_and_parameters4]

Flake rate in main: 33.33% (Passed 132 times, Failed 66 times)

Stack Traces | 0.004s run time

parameter_dataclass = <class 'sbi.inference.posteriors.posterior_parameters.VIPosteriorParameters'>
init_target_class = <class 'sbi.inference.posteriors.vi_posterior.VIPosterior'>
skipped_fields_and_parameters = {'device', 'potential_fn', 'prior', 'self', 'theta_transform', 'x_shape'}

    @pytest.mark.parametrize(
        ("parameter_dataclass", "init_target_class", "skipped_fields_and_parameters"),
        [
            (
                DirectPosteriorParameters,
                DirectPosterior,
                {"posterior_estimator", "prior", "device"},
            ),
            (
                ImportanceSamplingPosteriorParameters,
                ImportanceSamplingPosterior,
                {"potential_fn", "proposal", "device"},
            ),
            (
                MCMCPosteriorParameters,
                MCMCPosterior,
                {
                    "potential_fn",
                    "proposal",
                    "device",
                    "theta_transform",
                    "init_strategy_num_candidates",
                },
            ),
            (
                RejectionPosteriorParameters,
                RejectionPosterior,
                {"potential_fn", "device", "proposal"},
            ),
            (
                VIPosteriorParameters,
                VIPosterior,
                {"potential_fn", "prior", "theta_transform", "device"},
            ),
            (
                VectorFieldPosteriorParameters,
                VectorFieldPosterior,
                {
                    "vector_field_estimator",
                    "device",
                    "prior",
                    "sample_with",
                    "iid_method",
                    "iid_params",
                    "neural_ode_backend",
                    "neural_ode_kwargs",
                },
            ),
            (
                VectorFieldPosteriorParameters,
                VectorFieldBasedPotential,
                {
                    "vector_field_estimator",
                    "device",
                    "prior",
                    "x_o",
                    "enable_transform",
                    "max_sampling_batch_size",
                },
            ),
        ],
    )
    def test_signature_consistency(
        parameter_dataclass, init_target_class, skipped_fields_and_parameters
    ):
        """
        Test that the constructor (__init__) signature of a target class matches the
        signature of a corresponding parameter dataclass.
    
        This function compares the argument names, default values, and type annotations
        between the dataclass and the target class __init__ method, ignoring specified
        parameters passed in `skipped_fields_and_parameters`.
    
        Args:
            parameter_dataclass: The dataclass whose signature is used as reference.
            init_target_class: The class whose __init__ method signature is compared.
            skipped_fields_and_parameters (set): A set of parameter names to ignore during
                comparison (e.g., 'self', or fields not relevant for matching).
    
        Raises:
            AssertionError: If there is any mismatch in parameter names, default values,
                or type annotations between the dataclass and the class constructor.
        """
        dataclass_signature = inspect.signature(parameter_dataclass)
        class_signature = inspect.signature(init_target_class.__init__)
    
        skipped_fields_and_parameters.add("self")
        skipped_fields_and_parameters.add("x_shape")
    
        class_dict = {
            name: param
            for name, param in class_signature.parameters.items()
            if name not in skipped_fields_and_parameters
            and param.kind != inspect.Parameter.VAR_KEYWORD
        }
    
        dataclass_dict = {
            name: param
            for name, param in dataclass_signature.parameters.items()
            if name not in skipped_fields_and_parameters
        }
    
        # Compare if the dataclass and posterior_class have the same argument names
>       assert class_dict.keys() == dataclass_dict.keys(), (
            f"Parameter mismatch:\n"
            f"In class but not dataclass: {class_dict.keys() - dataclass_dict.keys()}\n"
            f"In dataclass but not class: {dataclass_dict.keys() - class_dict.keys()}"
        )
E       AssertionError: Parameter mismatch:
E         In class but not dataclass: {'parameters', 'modules'}
E         In dataclass but not class: {'hidden_features', 'z_score_theta', 'z_score_x', 'num_transforms'}
E       assert dict_keys(['q...', 'modules']) == dict_keys(['q... 'z_score_x'])
E         
E         Full diff:
E         - dict_keys(['q', 'vi_method', 'num_transforms', 'hidden_features', 'z_score_theta', 'z_score_x'])
E         + dict_keys(['q', 'vi_method', 'parameters', 'modules'])

tests/posterior_parameters_test.py:148: AssertionError

tests/save_and_load_test.py::test_picklability[NRE_B-VIPosteriorParameters]

Flake rate in main: 35.05% (Passed 63 times, Failed 34 times)

Stack Traces | 0.047s run time

inference_method = <class 'sbi.inference.trainers.nre.nre_b.NRE_B'>
posterior_parameters = <class 'sbi.inference.posteriors.posterior_parameters.VIPosteriorParameters'>
tmp_path = PosixPath('.../pytest-0/popen-gw1/test_picklability_NRE_B_VIPost0')
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')

    @pytest.mark.parametrize(
        "inference_method, posterior_parameters",
        (
            (NPE, DirectPosteriorParameters),
            (NPSE, VectorFieldPosteriorParameters),
            (FMPE, VectorFieldPosteriorParameters),
            pytest.param(NLE, MCMCPosteriorParameters, marks=pytest.mark.mcmc),
            pytest.param(NRE, MCMCPosteriorParameters, marks=pytest.mark.mcmc),
            pytest.param(NRE, VIPosteriorParameters, marks=pytest.mark.mcmc),
            (NRE, RejectionPosteriorParameters),
        ),
    )
    def test_picklability(
        inference_method,
        posterior_parameters,
        tmp_path,
        mcmc_params_fast: MCMCPosteriorParameters,
    ):
        num_dim = 2
        prior = utils.BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))
        x_o = torch.zeros(1, num_dim)
    
        theta = prior.sample((500,))
        x = theta + 1.0 + torch.randn_like(theta) * 0.1
    
        inference = inference_method(prior=prior)
        _ = inference.append_simulations(theta, x).train(max_num_epochs=1)
        if posterior_parameters is MCMCPosteriorParameters:
            posterior = inference.build_posterior(
                posterior_parameters=mcmc_params_fast
            ).set_default_x(x_o)
        else:
>           posterior = inference.build_posterior(
                posterior_parameters=posterior_parameters()
            ).set_default_x(x_o)

tests/save_and_load_test.py:53: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
.../trainers/nre/nre_base.py:360: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nre.nre_b.NRE_B object at 0x7f8c8cccb1c0>
estimator = RatioEstimator(
  (net): ResidualNet(
    (initial_layer): Linear(in_features=4, out_features=50, bias=True)
    (bloc...Standardize()
    (1): Identity()
  )
  (embedding_net_x): Sequential(
    (0): Standardize()
    (1): Identity()
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/sbc_test.py::test_running_sbc[NLE_A-vi-boxuniform-marginals]

Flake rate in main: 35.05% (Passed 63 times, Failed 34 times)

Stack Traces | 0.09s run time

method = <class 'sbi.inference.trainers.nle.nle_a.NLE_A'>
prior_type = 'boxuniform', reduce_fn_str = 'marginals', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')

    @pytest.mark.parametrize("reduce_fn_str", ("marginals", "posterior_log_prob"))
    @pytest.mark.parametrize("prior_type", ("boxuniform", "independent"))
    @pytest.mark.parametrize(
        "method, sampler",
        (
            (NPE, None),
            pytest.param(NLE, "mcmc", marks=pytest.mark.mcmc),
            pytest.param(NLE, "vi", marks=pytest.mark.mcmc),
            (NPSE, None),
        ),
    )
    def test_running_sbc(
        method,
        prior_type: str,
        reduce_fn_str: str,
        sampler: Optional[str],
        mcmc_params_fast: MCMCPosteriorParameters,
    ):
        """Test running inference and then SBC and obtaining nltp with different methods."""
        # Setup
        num_dim = 2
        if prior_type == "boxuniform":
            prior = BoxUniform(-torch.ones(num_dim), torch.ones(num_dim))
        else:
            prior = MultipleIndependent([
                Uniform(-torch.ones(1), torch.ones(1)) for _ in range(num_dim)
            ])
    
        # Test parameters
        num_simulations = 100
        max_num_epochs = 1
        num_sbc_runs = 2
        num_posterior_samples = 20
    
        likelihood_shift = -1.0 * ones(num_dim)
        likelihood_cov = 0.3 * eye(num_dim)
    
        # Helper function to simulate data
        def simulator(theta):
            return linear_gaussian(theta, likelihood_shift, likelihood_cov)
    
        # Build posterior
        posterior_kwargs = {}
        if method == NLE:
            posterior_kwargs = {
                "posterior_parameters": mcmc_params_fast
                if sampler == "mcmc"
                else VIPosteriorParameters()
            }
    
>       posterior = train_inference_method(
            method,
            prior,
            simulator,
            num_simulations=num_simulations,
            max_num_epochs=max_num_epochs,
            **posterior_kwargs,
        )

tests/sbc_test.py:118: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/sbc_test.py:63: in train_inference_method
    posterior = inferer.build_posterior(**kwargs)
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.nle_a.NLE_A object at 0x7f0fdafb2140>
estimator = NFlowsFlow(
  (net): Flow(
    (_transform): CompositeTransform(
      (_transforms): ModuleList(
        (0): Pointwi...ibution): StandardNormal()
    (_embedding_net): Sequential(
      (0): Standardize()
      (1): Identity()
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/sbc_test.py::test_running_sbc[NLE_A-vi-boxuniform-posterior_log_prob]

Flake rate in main: 35.05% (Passed 63 times, Failed 34 times)

Stack Traces | 0.089s run time

method = <class 'sbi.inference.trainers.nle.nle_a.NLE_A'>
prior_type = 'boxuniform', reduce_fn_str = 'posterior_log_prob', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')

    @pytest.mark.parametrize("reduce_fn_str", ("marginals", "posterior_log_prob"))
    @pytest.mark.parametrize("prior_type", ("boxuniform", "independent"))
    @pytest.mark.parametrize(
        "method, sampler",
        (
            (NPE, None),
            pytest.param(NLE, "mcmc", marks=pytest.mark.mcmc),
            pytest.param(NLE, "vi", marks=pytest.mark.mcmc),
            (NPSE, None),
        ),
    )
    def test_running_sbc(
        method,
        prior_type: str,
        reduce_fn_str: str,
        sampler: Optional[str],
        mcmc_params_fast: MCMCPosteriorParameters,
    ):
        """Test running inference and then SBC and obtaining nltp with different methods."""
        # Setup
        num_dim = 2
        if prior_type == "boxuniform":
            prior = BoxUniform(-torch.ones(num_dim), torch.ones(num_dim))
        else:
            prior = MultipleIndependent([
                Uniform(-torch.ones(1), torch.ones(1)) for _ in range(num_dim)
            ])
    
        # Test parameters
        num_simulations = 100
        max_num_epochs = 1
        num_sbc_runs = 2
        num_posterior_samples = 20
    
        likelihood_shift = -1.0 * ones(num_dim)
        likelihood_cov = 0.3 * eye(num_dim)
    
        # Helper function to simulate data
        def simulator(theta):
            return linear_gaussian(theta, likelihood_shift, likelihood_cov)
    
        # Build posterior
        posterior_kwargs = {}
        if method == NLE:
            posterior_kwargs = {
                "posterior_parameters": mcmc_params_fast
                if sampler == "mcmc"
                else VIPosteriorParameters()
            }
    
>       posterior = train_inference_method(
            method,
            prior,
            simulator,
            num_simulations=num_simulations,
            max_num_epochs=max_num_epochs,
            **posterior_kwargs,
        )

tests/sbc_test.py:118: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/sbc_test.py:63: in train_inference_method
    posterior = inferer.build_posterior(**kwargs)
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.nle_a.NLE_A object at 0x7f0fe47d7cd0>
estimator = NFlowsFlow(
  (net): Flow(
    (_transform): CompositeTransform(
      (_transforms): ModuleList(
        (0): Pointwi...ibution): StandardNormal()
    (_embedding_net): Sequential(
      (0): Standardize()
      (1): Identity()
    )
  )
)
prior = BoxUniform(Uniform(low: torch.Size([2]), high: torch.Size([2])), 1)
sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/sbc_test.py::test_running_sbc[NLE_A-vi-independent-marginals]

Flake rate in main: 35.05% (Passed 63 times, Failed 34 times)

Stack Traces | 0.085s run time

method = <class 'sbi.inference.trainers.nle.nle_a.NLE_A'>
prior_type = 'independent', reduce_fn_str = 'marginals', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')

    @pytest.mark.parametrize("reduce_fn_str", ("marginals", "posterior_log_prob"))
    @pytest.mark.parametrize("prior_type", ("boxuniform", "independent"))
    @pytest.mark.parametrize(
        "method, sampler",
        (
            (NPE, None),
            pytest.param(NLE, "mcmc", marks=pytest.mark.mcmc),
            pytest.param(NLE, "vi", marks=pytest.mark.mcmc),
            (NPSE, None),
        ),
    )
    def test_running_sbc(
        method,
        prior_type: str,
        reduce_fn_str: str,
        sampler: Optional[str],
        mcmc_params_fast: MCMCPosteriorParameters,
    ):
        """Test running inference and then SBC and obtaining nltp with different methods."""
        # Setup
        num_dim = 2
        if prior_type == "boxuniform":
            prior = BoxUniform(-torch.ones(num_dim), torch.ones(num_dim))
        else:
            prior = MultipleIndependent([
                Uniform(-torch.ones(1), torch.ones(1)) for _ in range(num_dim)
            ])
    
        # Test parameters
        num_simulations = 100
        max_num_epochs = 1
        num_sbc_runs = 2
        num_posterior_samples = 20
    
        likelihood_shift = -1.0 * ones(num_dim)
        likelihood_cov = 0.3 * eye(num_dim)
    
        # Helper function to simulate data
        def simulator(theta):
            return linear_gaussian(theta, likelihood_shift, likelihood_cov)
    
        # Build posterior
        posterior_kwargs = {}
        if method == NLE:
            posterior_kwargs = {
                "posterior_parameters": mcmc_params_fast
                if sampler == "mcmc"
                else VIPosteriorParameters()
            }
    
>       posterior = train_inference_method(
            method,
            prior,
            simulator,
            num_simulations=num_simulations,
            max_num_epochs=max_num_epochs,
            **posterior_kwargs,
        )

tests/sbc_test.py:118: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/sbc_test.py:63: in train_inference_method
    posterior = inferer.build_posterior(**kwargs)
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.nle_a.NLE_A object at 0x7f6fac20d720>
estimator = NFlowsFlow(
  (net): Flow(
    (_transform): CompositeTransform(
      (_transforms): ModuleList(
        (0): Pointwi...ibution): StandardNormal()
    (_embedding_net): Sequential(
      (0): Standardize()
      (1): Identity()
    )
  )
)
prior = MultipleIndependent(), sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/sbc_test.py::test_running_sbc[NLE_A-vi-independent-posterior_log_prob]

Flake rate in main: 35.05% (Passed 63 times, Failed 34 times)

Stack Traces | 0.087s run time

method = <class 'sbi.inference.trainers.nle.nle_a.NLE_A'>
prior_type = 'independent', reduce_fn_str = 'posterior_log_prob', sampler = 'vi'
mcmc_params_fast = MCMCPosteriorParameters(method='slice_np_vectorized', thin=1, warmup_steps=1, num_chains=1, init_strategy='resample', init_strategy_parameters=None, num_workers=1, mp_context='spawn')

    @pytest.mark.parametrize("reduce_fn_str", ("marginals", "posterior_log_prob"))
    @pytest.mark.parametrize("prior_type", ("boxuniform", "independent"))
    @pytest.mark.parametrize(
        "method, sampler",
        (
            (NPE, None),
            pytest.param(NLE, "mcmc", marks=pytest.mark.mcmc),
            pytest.param(NLE, "vi", marks=pytest.mark.mcmc),
            (NPSE, None),
        ),
    )
    def test_running_sbc(
        method,
        prior_type: str,
        reduce_fn_str: str,
        sampler: Optional[str],
        mcmc_params_fast: MCMCPosteriorParameters,
    ):
        """Test running inference and then SBC and obtaining nltp with different methods."""
        # Setup
        num_dim = 2
        if prior_type == "boxuniform":
            prior = BoxUniform(-torch.ones(num_dim), torch.ones(num_dim))
        else:
            prior = MultipleIndependent([
                Uniform(-torch.ones(1), torch.ones(1)) for _ in range(num_dim)
            ])
    
        # Test parameters
        num_simulations = 100
        max_num_epochs = 1
        num_sbc_runs = 2
        num_posterior_samples = 20
    
        likelihood_shift = -1.0 * ones(num_dim)
        likelihood_cov = 0.3 * eye(num_dim)
    
        # Helper function to simulate data
        def simulator(theta):
            return linear_gaussian(theta, likelihood_shift, likelihood_cov)
    
        # Build posterior
        posterior_kwargs = {}
        if method == NLE:
            posterior_kwargs = {
                "posterior_parameters": mcmc_params_fast
                if sampler == "mcmc"
                else VIPosteriorParameters()
            }
    
>       posterior = train_inference_method(
            method,
            prior,
            simulator,
            num_simulations=num_simulations,
            max_num_epochs=max_num_epochs,
            **posterior_kwargs,
        )

tests/sbc_test.py:118: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/sbc_test.py:63: in train_inference_method
    posterior = inferer.build_posterior(**kwargs)
.../trainers/nle/nle_base.py:291: in build_posterior
    return super().build_posterior(
.../inference/trainers/base.py:507: in build_posterior
    self._posterior = self._create_posterior(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <sbi.inference.trainers.nle.nle_a.NLE_A object at 0x7eff5906e980>
estimator = NFlowsFlow(
  (net): Flow(
    (_transform): CompositeTransform(
      (_transforms): ModuleList(
        (0): Pointwi...ibution): StandardNormal()
    (_embedding_net): Sequential(
      (0): Standardize()
      (1): Identity()
    )
  )
)
prior = MultipleIndependent(), sample_with = 'mcmc', device = 'cpu'
posterior_parameters = VIPosteriorParameters(q='maf', vi_method='rKL', num_transforms=5, hidden_features=50, z_score_theta='independent', z_score_x='independent')

    def _create_posterior(
        self,
        estimator: ConditionalEstimator,
        prior: Distribution,
        sample_with: Literal[
            "mcmc", "rejection", "vi", "importance", "direct", "sde", "ode"
        ],
        device: Union[str, torch.device],
        posterior_parameters: PosteriorParameters,
    ) -> NeuralPosterior:
        """
        Create a posterior object using the specified inference method.
    
        Depending on the value of `sample_with`, this method instantiates one of the
        supported posterior inference strategies.
    
        Args:
            estimator: The estimator that the posterior is based on.
            prior: A probability distribution that expresses prior knowledge about the
                parameters, e.g. which ranges are meaningful for them. Must be a PyTorch
                distribution, see FAQ for details on how to use custom distributions.
            sample_with: The inference method to use. Must be one of:
                - "mcmc"
                - "rejection"
                - "vi"
                - "importance"
                - "direct"
                - "sde"
                - "ode"
            device: torch device on which to train the neural net and on which to
                perform all posterior operations, e.g. gpu or cpu.
            posterior_parameters: Configuration passed to the init method for the
                posterior. Must be of type PosteriorParameters.
    
        Returns:
            NeuralPosterior object.
        """
    
        if isinstance(posterior_parameters, DirectPosteriorParameters):
            posterior_estimator = estimator
            if not isinstance(posterior_estimator, ConditionalDensityEstimator):
                raise TypeError(
                    f"Expected posterior_estimator to be an instance of "
                    " ConditionalDensityEstimator, "
                    f"but got {type(posterior_estimator).__name__} instead."
                )
            posterior = DirectPosterior(
                posterior_estimator=posterior_estimator,
                prior=prior,
                device=device,
                **asdict(posterior_parameters),
            )
        elif isinstance(posterior_parameters, VectorFieldPosteriorParameters):
            vector_field_estimator = estimator
            if not isinstance(vector_field_estimator, ConditionalVectorFieldEstimator):
                raise TypeError(
                    f"Expected vector_field_estimator to be an instance of "
                    " ConditionalVectorFieldEstimator, "
                    f"but got {type(vector_field_estimator).__name__} instead."
                )
            if sample_with not in ("ode", "sde"):
                raise ValueError(
                    "`sample_with` must be either",
                    f" 'ode' or 'sde', got '{sample_with}'",
                )
            posterior = VectorFieldPosterior(
                vector_field_estimator=vector_field_estimator,
                prior=prior,
                device=device,
                sample_with=sample_with,
                **asdict(posterior_parameters),
            )
        else:
            # Posteriors requiring potential_fn and theta_transform
            potential_fn, theta_transform = self._get_potential_function(
                prior, estimator
            )
            if isinstance(posterior_parameters, MCMCPosteriorParameters):
                posterior = MCMCPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, RejectionPosteriorParameters):
                posterior = RejectionPosterior(
                    potential_fn=potential_fn,
                    proposal=prior,
                    device=device,
                    **asdict(posterior_parameters),
                )
            elif isinstance(posterior_parameters, VIPosteriorParameters):
>               posterior = VIPosterior(
                    potential_fn=potential_fn,
                    theta_transform=theta_transform,
                    prior=prior,
                    device=device,
                    **asdict(posterior_parameters),
E                   TypeError: VIPosterior.__init__() got an unexpected keyword argument 'num_transforms'

.../inference/trainers/base.py:899: TypeError

tests/torchutils_test.py::TorchUtilsTest::test_searchsorted

Flake rate in main: 48.45% (Passed 50 times, Failed 47 times)

Stack Traces | 0.004s run time

.venv/lib/python3.10....../site-packages/xdist/remote.py:289: in pytest_runtest_logreport
    self.sendevent("testreport", data=data)
.venv/lib/python3.10....../site-packages/xdist/remote.py:126: in sendevent
    self.channel.send((name, kwargs))
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:912: in send
    self.gateway._send(Message.CHANNEL_DATA, self.id, dumps_internal(item))
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1629: in dumps_internal
    return _Serializer().save(obj)  # type: ignore[return-value]
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1647: in save
    self._save(obj)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1667: in _save
    dispatch(self, obj)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1744: in save_tuple
    self._save(item)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1667: in _save
    dispatch(self, obj)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1740: in save_dict
    self._write_setitem(key, value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1734: in _write_setitem
    self._save(value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1667: in _save
    dispatch(self, obj)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1740: in save_dict
    self._write_setitem(key, value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1734: in _write_setitem
    self._save(value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1667: in _save
    dispatch(self, obj)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1740: in save_dict
    self._write_setitem(key, value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1734: in _write_setitem
    self._save(value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1667: in _save
    dispatch(self, obj)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1740: in save_dict
    self._write_setitem(key, value)
.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1734: in _write_setitem
    self._save(value)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = <execnet.gateway_base._Serializer object at 0x7f210c584e50>
obj = tensor([0.0000, 0.1111, 0.2222, 0.3333, 0.4444, 0.5556, 0.6667, 0.7778, 0.8889])

    def _save(self, obj: object) -> None:
        tp = type(obj)
        try:
            dispatch = self._dispatch[tp]
        except KeyError:
            methodname = "save_" + tp.__name__
            meth: Callable[[_Serializer, object], None] | None = getattr(
                self.__class__, methodname, None
            )
            if meth is None:
>               raise DumpError(f"can't serialize {tp}") from None
E               execnet.gateway_base.DumpError: can't serialize <class 'torch.Tensor'>

.venv/lib/python3.10....................................................../site-packages/execnet/gateway_base.py:1665: DumpError